# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/14 15:24 
# @Author  : shutao
# @FileName: lfw_dataset.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

from loguru import logger
import numpy as np
import cv2 as cv
import traceback
import random
import paddle
import os

from paddle.io import IterableDataset


def _create_datalist(root_path):
    """
    创建数据路径列表
    Args:
        root_path: 数据根目录

    Returns:
        'datalist': 返回数据列表
            数组第一级为类别编号，第二级为类内编号
            [[p1, p2, ..., pn], [], [], ..., []]

        'count': 总数据条数
    """
    datalist = []
    count = 0
    for label, _person in enumerate(os.listdir(root_path)):
        person_path = os.path.join(root_path, _person)
        datalist.append([])
        for idx, _image in enumerate(os.listdir(person_path)):
            image_path = os.path.join(person_path, _image)
            # print(image_path)
            datalist[label].append(image_path)
            count += 1
    return datalist, count


class FaceDataset(IterableDataset):
    """ 人脸匹配 lfw 数据集专用迭代器 """

    def __init__(self,
                 root_path="E:\\work\\Dataset\\lfw",
                 img_size=(224, 224)):
        super(FaceDataset, self).__init__()
        self.img_size = img_size

        self.datalist, self.num_samples = _create_datalist(root_path)
        self.label_num = len(self.datalist)  # 类数量

    def _process(self, img):
        return img

    def _read_img(self, path):
        img = cv.imread(path)
        img = self._process(img)
        img = cv.resize(img, self.img_size, interpolation=cv.INTER_LINEAR)
        img = np.transpose(img, (2, 0, 1))  # 转换数据维度
        img = (img - 128) / 255  # 归一化
        return img

    def __iter__(self):
        for i in range(self.num_samples):
            try:
                index = random.randint(0, self.label_num - 1)
                if i % 2 == 0:  # 相同人脸
                    while len(self.datalist[index]) < 2:  # 选择类内数量大于2的人脸类别
                        index = random.randint(0, self.label_num - 1)
                    rand_num1 = random.randint(0, len(self.datalist[index]) - 1)
                    img1 = self._read_img(self.datalist[index][rand_num1])

                    rand_num2 = random.randint(0, len(self.datalist[index]) - 1)
                    while rand_num2 == rand_num1:
                        rand_num2 = random.randint(0, len(self.datalist[index]) - 1)
                    img2 = self._read_img(self.datalist[index][rand_num2])
                    label = np.array([1])  # 为 1 代表相同人脸
                else:
                    while True:
                        _index = random.randint(0, self.label_num - 1)
                        if _index != index:
                            break
                    rand_num1 = random.randint(0, len(self.datalist[index]) - 1)
                    img1 = self._read_img(self.datalist[index][rand_num1])

                    rand_num2 = random.randint(0, len(self.datalist[_index]) - 1)
                    img2 = self._read_img(self.datalist[_index][rand_num2])
                    label = np.array([0])  # 为 0 代表不同人脸
                yield img1, img2, label
            except Exception as e:
                logger.error("数据读取错误，错误信息: \n{}".format(traceback.format_exc()))


if __name__ == '__main__':
    # dataset = FaceDataset()
    # dataloader = paddle.io.DataLoader(
    #     dataset,
    #     num_workers=0,
    #     batch_size=4,
    #     drop_last=True
    # )
    # for img1, img2, label in dataloader:
    #     print(img1.shape, img2.shape, label)
    pass

